#!/usr/bin/python
'''
Created by Colin Gillespie <c.gillespie@ncl.ac.uk>

Last Changed on: $Date: 2008-03-13 19:50:58 +0000 (Thu, 13 Mar 2008) $ 
Work in progress

'''

__revision__ = '$Revision: 73 $'

import sys, re
from swiginac import diff

try:
    import psyco
    psyco.full()
except ImportError:
    pass

#Import non-standard modules
from mrange import cumRange
from VariableDict import Variables
from GetMomentEquation import getMomentEquation
from GetLambdaDict import GetLambdaDict
from GetEvents import GetEvents
import Maple

def getHeader():
    header = '''#include <stdio.h>
#include <math.h>    
#include <gsl/gsl_errno.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_odeiv.h>
gsl_matrix *rk_out;
#define MSET(rk_out,i,j,k) gsl_matrix_set(rk_out,i,j,k)\n
'''
    return header

def getParametersDict(GDD):#, list_of_pars):
    par_dict = {}
    for pa in GDD.pa_dict.keys():
        par_dict[pa] = str(GDD.pa_dict[pa]['value'])

    return par_dict


def getfunc(GDD, reactions_dict):
    expression = 'int func(double t, const double kappa[], double dkappa[], void *params)\n'
    expression += '{\n'
    par_dict = getParametersDict(GDD)
    for value, key in par_dict.items():
        expression += '\tdouble %s=%s;\n' % (value, key)
    
    p=re.compile('(kappa\[\d*\])\*\*(\d*)',re.VERBOSE)
    cum_dict = Variables().cum_look_up
    dcum_dict = Variables().dcum_look_up
    for value, key in reactions_dict.items():
        reaction = p.sub(r'pow(\g<1>,\g<2>)', str(key))
        expression += '\t%s=%s;\n' %(dcum_dict[value], reaction)
        createJacobianDict(value, key)
    expression += '\treturn GSL_SUCCESS;\n}\n\n'
    return  expression



def createJacobianDict(dkappa=None, reaction=None, jac_dict={}):
    if not dkappa and not reaction:
        return jac_dict
    
    pos_i = str(Variables(dkappa).get())[6:-1]    
    max_values = Variables().getMaxValues()
    for i,values in enumerate(cumRange(max_values)):
        kappa= Variables(values).get()
        pos_j = str(kappa)[6:-1]
        position = (pos_i, pos_j)
        jac_dict[position] = str(diff(reaction, kappa))
        
def getjac(reactions_dict):
    l = len(reactions_dict)
    expression = 'int jac(double t, const double kappa[], double *dfdy,double dfdt[], void *params)\n{'
    expression += '\tgsl_matrix_view dfdy_mat = gsl_matrix_view_array (dfdy, %s, %s);\n' % (l, l)
    expression += '\tgsl_matrix *m = &dfdy_mat.matrix;\n'
    jac_dict = createJacobianDict()
    for value, key in jac_dict.items():
        expression += '\tgsl_matrix_set (m, %s, %s, %s);\n' % (value[0], value[1], str(key))

    for i in range(len(reactions_dict)):
        expression +='\tdfdt[%s]=0;\n' % str(i)
    
    expression += '\treturn GSL_SUCCESS;\n}'
    return expression




def getrk_main(reactions_dict, initial_conditions_dict):
    l = len(reactions_dict)
    expression = '''int rk_main(int max_time)
    {
        const gsl_odeiv_step_type *T = gsl_odeiv_step_rk4;
        gsl_odeiv_step *stp = gsl_odeiv_step_alloc (T, %s);
        gsl_odeiv_control *c = gsl_odeiv_control_y_new (1e-6, 0.0);/*error allowances*/
        gsl_odeiv_evolve * e = gsl_odeiv_evolve_alloc (%s);
        gsl_odeiv_system sys = {func, 0, %s, 0};
  
        double t = 0.0;
        double h = 1e-6;
        int i,j;  
  
  ''' % (l, l, l)
    initial_conditions_list = l*[0]
    for value, key in initial_conditions_dict.items():
        index = int(Variables().cum_look_up[value][6:-1])
        initial_conditions_list[index] = float(key)


    ic = str(initial_conditions_list)[1:-1]
    expression += 'double kappa[%s] = {%s};\n\n' % (l, ic)
    expression +='''for (i = 1; i <=(int) max_time; i++)
    {
     
      while (t < i)
	    {
	        gsl_odeiv_evolve_apply (e, c, stp, &sys, &t, i, &h, kappa);
    	}
    for(j=0;j<3;j++) {
	    MSET(rk_out,j,i-1,kappa[j]);
        }
    }


  gsl_odeiv_evolve_free (e);
  gsl_odeiv_control_free (c);
  gsl_odeiv_step_free (stp);
  return 0;
}\n
    '''
    return expression 

def getMain(reactions_dict):
    expression = '''int main()
    {
    int i,j;
    double max_time=100.0;
    rk_out=gsl_matrix_alloc(%s, (int) max_time);
  
    rk_main(max_time);
    for(i=0;i<max_time;i++){
        for(j=0;j<3;j++){
    '''% str(len(reactions_dict))
    expression += '''printf("%f ", gsl_matrix_get(rk_out,j,i));
        }
    printf("\\n"); 
    }
    return 0;
    }\n\n''' 
    return expression



def getDifferentialEquations(output_file='tmp.c'):
    '''Creates the maple file with moment equations.
    
    Arguments:
        sbml_file - sbml file name
        output_file - outout file
        moment_order - what moment order to go to.
    '''
    #sbml_file = '/home/ncsg3/moment_closure/trunk/xml_files/smfsb/ch07-lac-stoch.xml'
    sbml_file = '/tmp/bruss.xml'
    moment_order = 1
    Variables().setMomentOrder(moment_order)
    GDD = GetLambdaDict(sbml_file = sbml_file)
    lambda_dict = GDD.getLambdaDict()
    Variables().setMaxValues([moment_order] * len(GDD.empty_list))
    
    list_of_species = []
    reactions_dict = {}
    
    constant_sp_dict = {}
    #Output the species
    initial_conditions_dict = {}
    for values in cumRange(Variables().getMaxValues()):
        if sum(values) > 1: #higher order cumulant
            initial_conditions_dict[values] = str(0.0)
        else:
            index = list(values).index(1)
            sp = GDD.m.getSpecies(index)
            initial_conditions_dict[values] = str(sp.getInitialAmount())


        mom_eq = getMomentEquation(lambda_dict = lambda_dict, cumulant_equation = values)
        reactions_dict[values] = mom_eq
        if mom_eq == 0:
            reactions_dict[values] = mom_eq
            if sum(values) > 1: #higher order cumulant
                constant_sp_dict[str(Variables(tuple(values), True).get())[:-3]] = 0
            else:
                index = list(values).index(1)
                sp = GDD.m.getSpecies(index)
                constant_sp_dict[str(Variables(tuple(values), True).get())[:-3]] = sp.getInitialAmount()
          
    print getHeader()
    print getfunc(GDD, reactions_dict)
    print getjac(reactions_dict)
    print getrk_main(reactions_dict, initial_conditions_dict)
    print getMain(reactions_dict)
    return True

if __name__ == '__main__':
    getDifferentialEquations()











